
#include <vector>
#include "AESSymmetricKey.h"
#include "ECDSAPrivateKey.h"
#include "DSAPrivateKey.h"
#include "RSAPrivateKey.h"
#include "RSAPublicKey.h"
#include "CryptOChain.h"
#include "EncryptWrapper.h"

/** General initialisation function for the encryption wrapper.
 * 
 * @param enckey EncryptionKey to use for encrypting data if present
 * and not NULL.
 * 
 * @param sgnkey SigningKey to use to sign the data stream if present
 * and not NULL.
 * 
 * @param compress A flag to specify an additional compression step if
 * true. */
void EncryptWrapper::initEncryptWrapper(EncryptionKey *enckey, SigningKey *sgnkey, int compress)
{
   if (enckey)
   {
      switch (enckey->getKeyType())
      {
         case AbstractKey::KEY_TYPE_PUBLIC_RSA:
            { enckeycopy = new RSAPublicKey(*((RSAPublicKey *)enckey)); break; }
         case AbstractKey::KEY_TYPE_SYMM_AES:
            { enckeycopy = new AESSymmetricKey(*((AESSymmetricKey *)enckey)); break; }
         default:
            { errprint("unknown encryption key type"); enckeycopy = NULL; break; }
      }
   } else
   {
      errprint("encrypt wrapper has no key: can't do encryption!");
      enckeycopy = NULL;
   }
   if (sgnkey)
   {
      switch (sgnkey->getKeyType())
      {
         case AbstractKey::KEY_TYPE_PRIVATE_RSA:
            { sgnkeycopy = new RSAPrivateKey(*((RSAPrivateKey *)sgnkey)); break; }
         case AbstractKey::KEY_TYPE_PRIVATE_DSA:
            { sgnkeycopy = new DSAPrivateKey(*((DSAPrivateKey *)sgnkey)); break; }
         case AbstractKey::KEY_TYPE_PRIVATE_ECDSA:
            { sgnkeycopy = new ECDSAPrivateKey(*((ECDSAPrivateKey *)sgnkey)); break; }
         default:
            { errprint("unknown signing key type"); sgnkeycopy = NULL; break; }
      }
   } else
   {
      sgnkeycopy = NULL;
   }
   docompress = (compress != 0);
   trusted = 0;
   apr_status_t status;
   mp = NULL;
   status = apr_pool_create(&mp, NULL);
   if (mp == NULL)
   {
      errprint("failed to create memory pool");
   }
   CheckAPRError(status);
}

/** Constructs an encryption process which encrypts with the given
 * EncryptionKey.
 * 
 * @param enckey EncryptionKey to use for encrypting data if present
 * and not NULL. */
EncryptWrapper::EncryptWrapper(EncryptionKey *enckey)
{
   initEncryptWrapper(enckey, NULL, FALSE);
}

/** Constructs an encryption process which encrypts the data with the
 * given EncryptionKey and subsequently signs the result with the
 * given SigningKey.
 * 
 * @param enckey EncryptionKey to use for encrypting data if present
 * and not NULL.
 * 
 * @param sgnkey SigningKey to use to sign the data stream if present
 * and not NULL. */
EncryptWrapper::EncryptWrapper(EncryptionKey *enckey, SigningKey *sgnkey)
{
   initEncryptWrapper(enckey, sgnkey, FALSE);
}

/** Constructs an encryption process which can optionally apply a
 * compression step, before encrypting with the given EncryptionKey.
 * 
 * @param enckey EncryptionKey to use for encrypting data if present
 * and not NULL.
 * 
 * @param docompress A flag to specify an additional compression step
 * if true. */
EncryptWrapper::EncryptWrapper(EncryptionKey *enckey, int docompress)
{
   initEncryptWrapper(enckey, NULL, docompress);
}

/** Constructs an encryption process which can optionally apply a
 * compression step, before encrypting with the given EncryptionKey
 * and signing the result with the given SigningKey.
 * 
 * @param enckey EncryptionKey to use for encrypting data if present
 * and not NULL.
 * 
 * @param sgnkey SigningKey to use to sign the data stream if present
 * and not NULL.
 * 
 * @param docompress A flag to specify an additional compression step
 * if true. */
EncryptWrapper::EncryptWrapper(EncryptionKey *enckey, SigningKey *sgnkey, int docompress)
{
   initEncryptWrapper(enckey, sgnkey, docompress);
}

/** Destructor for the encryption wrapper. */
EncryptWrapper::~EncryptWrapper()
{
   if (enckeycopy) delete enckeycopy;
   if (sgnkeycopy) delete sgnkeycopy;
   apr_pool_clear(mp);
   apr_pool_destroy(mp);
}

/** Apply the encryption process to a buffer of data, returning the
 * encoded data.
 * 
 * @param buffer Buffer for uncoded data input.
 * 
 * @param length Length of buffer for data input.
 * 
 * @param[out] newlength Length of output buffer
 * 
 * @return A char buffer into which the encoded data is put. The
 * caller is responsible for deleting the new buffer instance. */
char *EncryptWrapper::encryptDataBuffer(const char *buffer, apr_size_t length, apr_size_t *newlength)
{
   CryptOChain *crypty = new CryptOChain(mp);
   crypty->buildEncryptionChain(enckeycopy, sgnkeycopy, docompress);
   char *newdata = crypty->processDataBlock(buffer, length, newlength);
   delete crypty;
   return newdata;
}

/** Apply the encryption process to a pipe of data in the form of an
 * APRBucketBrigade, returning a pointer to a newly created
 * APRBucketBrigade containing the encoded data.
 * 
 * @param bbrgde The bucket brigade input of uncoded data.
 * 
 * @return New bucket brigade representing the complete encoded data
 * stream. The caller is responsible for deleting the APRBucketBrigade
 * instance. */
APRBucketBrigade *EncryptWrapper::encryptToBucket(APRBucketBrigade *bbrgde)
{
   APRBucketBrigade *outputbrgde = new APRBucketBrigade(mp);
   CryptOChain *crypty = new CryptOChain(*bbrgde, *outputbrgde, mp);
   crypty->buildEncryptionChain(enckeycopy, sgnkeycopy, docompress);
   crypty->processPipeBegin();
   int prociter = 1;
   // We may need to disallow finishing before a isClosed
   while (prociter)
      prociter = crypty->processPipeStep(ENV_READ_BUFFER, ENV_READ_BUFFER);
   crypty->processPipeEnd();
   trusted = crypty->processWasSecure();
   delete crypty;
   if (prociter == -1)
   {
      delete outputbrgde;
      return NULL;
   }
   return outputbrgde;
}
